#!/usr/bin/env python3
# -*- coding: utf-8 -*
import os
import re
import json
import copy
import numpy as np
import sys 
sys.path.append("../..")
from package.utility import find_targets
from package.utility import ensure_exists

def InputdirPath(topdir): #输入一个算例文件夹名，输出一个存着文件夹地址的list
    "遍历文件夹得到文件夹地址"
    topdir = os.path.abspath(topdir)
    ensure_exists(topdir)

    print(f">> scan runcase under {topdir}") 
    runcases = find_targets(topdir, "srun_out.log", suffix="-xtype f") 
    runcase_num = len(runcases)
    print(f">> {runcase_num} runcases found!")
    return runcases

def writedict2json(dict):
    print("\nwrite data into:")
    for i in dict:
        dir = dict[i].pop("dir")
        jsondata = json.dumps(dict[i], sort_keys=True, indent=4, separators=(',', ': '))
        with open(f"{dir}/{i}-compute_length.json","w") as fo:
            fo.write(jsondata)
        print(f"  {dir}/{i}-compute_length.json")
    return

def getcasepath(casename, mypath, dict):
    #获取casename的文件夹名字
    pathlist = mypath.split("/")
    dirstr = "/"
    for j in pathlist:
        if j != casename:
            dirstr = os.path.join(dirstr, j)
        else:
            break
    dirstr = os.path.join(dirstr, casename)
    dict[casename]["dir"] = dirstr
    return dict

def main(mydir):
    Inpttdirlist = InputdirPath(mydir)
    compute_length_dict = {}
    for i in Inpttdirlist :
        filepath = os.path.dirname(i)
        with open(f"{filepath}/srun_config.json", 'r') as fo:
            jsonData = fo.read()
            text = json.loads(jsonData)
            npkey = "np_" + str(text["proc_num"])
            casename = text["testcase"]
            if casename not in compute_length_dict:
                compute_length_dict[casename] = {}
                compute_length_dict = getcasepath(casename, filepath, compute_length_dict)
            if npkey not in compute_length_dict[casename]:
                compute_length_dict[casename][npkey] = {}
        with open(i, 'r') as fo:
            loglines = fo.readlines()
            for logline in loglines:
                # >> level 6, communicate length(max,ave): (125,62.226562)
                search_result = re.match(r" >> level (\d+), communicate length\(max,ave\): \((\d+),(\d+.\d+)\)", logline, re.M|re.I)
                if(search_result):
                    key_level = "level_" + search_result[1]
                    compute_length_dict[casename][npkey][key_level] = {}
                    compute_length_dict[casename][npkey][key_level]["max"] = search_result[2]
                    compute_length_dict[casename][npkey][key_level]["ave"] = search_result[3]
                    # print(search_result[0], search_result[1], search_result[2], search_result[3])
    # print(compute_length_dict)
    writedict2json(compute_length_dict)

if __name__=="__main__":
    mydir = "/home/mrz/data/23_2_23_bottleneckresult_huang/23_3_6_compute_length/Out"
    main(mydir)